# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import sys
import collections
import json
import pandas as pd

from absl import app, flags

# matplotlib.use('Agg')
import matplotlib.pyplot as plt
FLAGS = flags.FLAGS


def load_result(filename):
  with open(filename) as infile:
    data = json.load(infile)
  return data


def get_pnl(data):
  ts = []
  pnl_list = collections.defaultdict(list)
  pos_list = collections.defaultdict(list)
  for pnl in data:
    ts.append(pnl['timestamp'])
    for product, info in pnl['data'].items():
      pnl_list[product].append(info['pnl'])
      pos_list[product].append(info['pos'])

  ts = pd.to_datetime(ts)
  df = pd.DataFrame(pnl_list)
  df = df.set_index(ts)
  pnl_df = df

  ts = pd.to_datetime(ts)
  df = pd.DataFrame(pos_list)
  df = df.set_index(ts)
  pos_df = df
  return pnl_df, pos_df


def main(_):
  data = load_result(sys.argv[1])
  df, pos_df = get_pnl(data)
  for n in range(3):
    plt.subplot(3, 1, n + 1)
    plt.plot(df.index, df.iloc[:, n], label=df.columns[n])
    plt.legend()
  plt.show()

  for n in range(3):
    plt.subplot(3, 1, n + 1)
    plt.plot(pos_df.index, pos_df.iloc[:, n], label=pos_df.columns[n])
    plt.legend()
  plt.show()


if __name__ == '__main__':
  flags.DEFINE_string(
      'symbol',
      None,
      'Specify product full symbol.',
  )

  flags.DEFINE_string('trading_date', None, 'Trading date.')

  flags.DEFINE_string('home_dir', None, 'Home directory.')
  app.run(main)
